#![allow(unused_imports)]
use anyhow::{Context, Result};
use bytes::Bytes;
use eztrade_dto::{
    DataTransferObject, Message,
    tick::{Orderstack, TickData},
    stage, user::{Authority}, events::{self, SubscribeEvent, GlobalEvent, Frequency, DataType}
};
use tracing::Instrument;
use crate::core::sandbox::{SandboxTrait, create_sandbox};
use crate::utils::SANDBOX_CONFIG;
use salvo::prelude::*;
use salvo::proto::quic::BidiStream;
use salvo::proto::webtransport;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::pin;


async fn bidi_heartbeat<S>(mut stream: S) -> Result<()>
where
    S: Unpin + AsyncRead + AsyncWrite,
{
    log::debug!("Opening bidirectional stream");

    // 发送初始消息
    let hello = Message::Notification("Hello from server".to_string());
    if let Err(e) = stream
        .write_all(&hello.to_binary_bytes()?)
        .await
    {
        log::error!("Failed to send initial message: {e}");
        stream.shutdown().await?;
        return Err(e.into());
    }

    // 创建定时器，每2秒发送一次消息
    let mut ping_interval = tokio::time::interval(Duration::from_secs(2));
    // let mut resp = Vec::new();

    loop {
        let _ = ping_interval.tick().await;
        if let Err(e) = stream.write_all(b"ping").await {
            log::error!("Failed to send heartbeat: {e}");
            break;
        }
    }
    log::debug!("Closing bidirectional stream");
    stream.shutdown().await?;
    Ok(())
}

async fn di_heartbeat<T, R>(send: T, recv: R) -> Result<()>
where
    T: AsyncWrite,
    R: AsyncRead,
{
    pin!(send);
    pin!(recv);

    let mut buf = Vec::new();
    match recv.read_to_end(&mut buf).await {
        Ok(_) => log::trace!("Received {} bytes in echo stream", buf.len()),
        Err(e) => {
            log::error!("Failed to read from echo stream: {e}");
            // return Err(e.into());
        }
    }

    match buf.as_slice() {
        b"ping" => {
            buf = b"pong".to_vec();
            let response = Bytes::from(buf);
            match send.write_all(&response).await {
                Ok(_) => {}
                Err(e) => {
                    log::error!("Stream Failed to send pong: {e}");
                }
            }
        }
        b"pong" => {
            send.shutdown().await?;
            return Ok(());
        }
        _ => {
            log::debug!("Received unknown message, ignoring");
            return Ok(());
        }
    }

    Ok(())
}

#[handler]
#[tracing::instrument(skip_all, fields(stage_id = %req.param::<u64>("id").unwrap_or(0)))]
pub async fn connect(req: &mut Request, resp: &mut Response) -> Result<(), salvo::Error> {
    // 获取当前span上下文
    let span = tracing::Span::current();
    // 从请求参数中获取ID
    let stage_id = req.param::<u64>("id").unwrap();

    // 检查是否为 WebTransport 连接
    let session = match req.web_transport_mut().await {
        Ok(session) => session,
        Err(e) => {
            log::warn!("非 WebTransport 请求访问 /wt 端点: {}", e);
            // 返回 426 Upgrade Required 状态码，表示需要协议升级
            resp.status_code(StatusCode::UPGRADE_REQUIRED);
            resp.render("WebTransport connection required");
            return Ok(());
        }
    };

    let session_id = session.session_id();
    tracing::debug!("WebTransport 连接已建立, 会话ID: {:?}", session_id);

    // 该版本无法处理单向Datagram的发送(可接收), 所以只能处理单向Stream的接发
    // 服务器主动发起双向流
    let stream = session.open_bi(session_id).await?;    
    tokio::spawn(bidi_heartbeat(stream).instrument(span.clone()));
    
    // 创建或获取沙盒实例
    let stage_info = eztrade_dto::stage::StageInfo {
        id: stage_id,
        stage_type: eztrade_dto::stage::StageType::SIM,
        name: format!("Stage-{}", stage_id),
        description: format!("Auto-generated stage for ID {}", stage_id),
        phase: eztrade_dto::stage::StagePhase::Running,
        time_unit_us: SANDBOX_CONFIG.time_unit_us, // 500ms
        margin: None,
        fee: None,
        market_sim: Some(eztrade_dto::stage::MarketSimConfig {
            time_multiplier: SANDBOX_CONFIG.time_scale,
            liquidity_repair: 1.0,
            trade_lag_sec: 0.0,
            data_source: None,
            simulate_start_us: chrono::Utc::now().timestamp_micros() - 60 * 60 * 1_000_000,
            simulate_end_us: chrono::Utc::now().timestamp_micros(),
        }),
        created_at_us: chrono::Utc::now().timestamp_micros(),
    };
    
    let sandbox = create_sandbox(stage_id, stage_info)?;
    sandbox.start()?; // 启动沙盒环境
    // 全局订阅
    sandbox.subscribe_global(GlobalEvent::new(
        DataType::Tick,
        Some("ANY".to_string()),
        None,
    ))?;
    
    let mut option_receiver: Option<tokio::sync::broadcast::Receiver<Message>> = None;
    let mut option_auth: Option<Authority> = None;
    
    log::debug!("waiting for authority message or data receiver");
    // 阻塞循环等待鉴权
    loop {
        tokio::select! {
            // 接收单向流鉴权
            uni_stream = session.accept_uni() => {
                let (_id, mut stream) = uni_stream?.unwrap();
                let mut buf = Vec::new();
                stream.read_to_end(&mut buf).await?;
                match Message::from_binary_bytes(&Bytes::from(buf)){
                    Ok(message) => {
                        if option_receiver.is_none() { // 等待鉴权
                            if let Message::Authority(mut auth) = message {
                                // 从沙盒获取数据接收器
                                auth.generate()?;
                                let data_receiver = sandbox.get_receiver(auth.client_id.unwrap())?;
                                option_receiver.replace(data_receiver);
                                option_auth.replace(auth);
                                break;
                            }
                        }
                    }
                    _ => {}
                };
            }

            // 鉴权超时
            _ = tokio::time::sleep(Duration::from_secs(10)) => {
                if option_receiver.is_none() {
                    log::debug!("No active receiver, closing session");
                    return Ok(());
                }
            }

            else => {
                break
            }
        }
    }

    log::debug!("Received authority message, starting data receiver");
    let mut receiver = option_receiver.take().unwrap();
    let auth = option_auth.take().unwrap();
    let client_id = auth.client_id.unwrap();
    sandbox.subscribe(client_id, SubscribeEvent::new(
        DataType::OrderDone,
        client_id,
        None,
        None,
    ))?;
    // 反馈鉴权信息
    let response = Message::Authority(auth).to_binary_bytes().unwrap();
    let mut send = session.open_uni(session_id).await?;
    send.write_all(&response).await?;

    // 主循环处理各种连接类型
    loop {
        tokio::select! {
            
            // 接收客户端单向流 (主要处理Message对象)
            uni_stream = session.accept_uni() => {
                let (_id, mut stream) = uni_stream?.unwrap();
                let mut buf = Vec::new();
                if let Err(e) = stream.read_to_end(&mut buf).await {
                    log::error!("Failed to read from uni stream: {}", e);
                    continue;
                }
                match Message::from_binary_bytes(&Bytes::from(buf)){
                    Ok(message) => {
                        // 将消息发送给沙盒处理
                        let response = span.in_scope(|| sandbox.process_message(client_id, message.clone()));
                        if let Some(response) = response {
                            let mut send = session.open_uni(session_id).await?;
                            if let Err(e) = send.write_all(&response.to_binary_bytes()?).await {
                                log::error!("Failed to send response: {}", e);
                            }
                        }
                    }
                    Err(e) => {
                        span.in_scope(|| {
                            log::error!("Failed to deserialize message: {}", e);
                        });
                        continue;
                    }
                };
            }

            // 主动发送单向流的订阅数据
            result = receiver.recv() => {
                if let Ok(message) = result {
                    let message_bytes = message.to_binary_bytes().unwrap();
                    let mut send = session.open_uni(session_id).await?;
                    send.write_all(&message_bytes).await?;
                }
            }

            // 处理客户端发起的双向ping/pong流(额外发起的)
            stream = session.accept_bi() => {
                if let Some(webtransport::server::AcceptedBi::BidiStream(_, stream)) = stream? {
                    let (send, recv) = stream.split();
                    tokio::spawn(
                        di_heartbeat(send, recv).instrument(span.clone())
                    );
                }
            }

            else => {
                break
            }
        }
    }
    // never reach
    Ok(())
}
